// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/conv.h>

#include <pollux/vector/builder_type_utils.h>

namespace kumo::pollux {
    template<typename T>
    SequenceVector<T>::SequenceVector(
        pollux::memory::MemoryPool *pool,
        size_t length,
        VectorPtr sequenceValues,
        BufferPtr sequenceLengths,
        const SimpleVectorStats<T> &stats,
        std::optional<int32_t> distinctCount,
        std::optional<vector_size_t> nullCount,
        std::optional<bool> isSorted,
        std::optional<ByteCount> represented_bytes,
        std::optional<ByteCount> storageByteCount)
        : SimpleVector<T>(
              pool,
              sequenceValues->type(),
              VectorEncoding::Simple::SEQUENCE,
              BufferPtr(nullptr),
              length,
              stats,
              distinctCount,
              nullCount,
              isSorted,
              represented_bytes,
              storageByteCount),
          sequenceValues_(std::move(sequenceValues)),
          sequenceLengths_(std::move(sequenceLengths)) {
        POLLUX_CHECK(sequenceLengths_ != nullptr, "Sequence Lengths must be not null");
        POLLUX_CHECK_EQ(
            sequenceLengths_->size(),
            BaseVector::byteSize<SequenceLength>(sequenceValues_->size()),
            "Sequence Lengths must be sized to hold lengths for each sequence value");
        setInternalState();
    }

    template<typename T>
    void SequenceVector<T>::setInternalState() {
        if (sequenceValues_->is_scalar()) {
            scalarSequenceValues_ =
                    reinterpret_cast<SimpleVector<T> *>(sequenceValues_.get());
        }
        lengths_ = sequenceLengths_->as<vector_size_t>();
        lastIndexRangeEnd_ = lengths_[0];
        BaseVector::in_memory_bytes_ += sequenceValues_->in_memory_bytes();
        BaseVector::in_memory_bytes_ += sequenceLengths_->capacity();
    }

    template<typename T>
    bool SequenceVector<T>::isNullAtFast(vector_size_t idx) const {
        size_t offset = offsetOfIndex(idx);
        return sequenceValues_->is_null_at(offset);
    }

    template<typename T>
    const T SequenceVector<T>::valueAtFast(vector_size_t idx) const {
        size_t offset = offsetOfIndex(idx);
        return scalarSequenceValues_->value_at(offset);
    }

    template<typename T>
    std::unique_ptr<SimpleVector<uint64_t> > SequenceVector<T>::hash_all() const {
        POLLUX_CHECK(is_scalar(), "Complex types not yet supported");
        // TODO T70734527 dealing with zero length vector
        if (BaseVector::length_ == 0) {
            return nullptr;
        }
        // If there is at least one value, then indices_ is set and has a pool.
        auto sequenceCount = numSequences();
        BufferPtr hashes =
                AlignedBuffer::allocate<uint64_t>(sequenceCount, BaseVector::pool_);
        uint64_t *rawHashes = hashes->asMutable<uint64_t>();
        for (size_t i = 0; i < sequenceCount; ++i) {
            rawHashes[i] = sequenceValues_->hash_value_at(i);
        }
        auto hashValues = std::make_shared<FlatVector<uint64_t> >(
            BaseVector::pool_,
            BIGINT(),
            BufferPtr(nullptr),
            sequenceCount,
            hashes,
            std::vector<BufferPtr>(0) /*stringBuffers*/,
            SimpleVectorStats<uint64_t>{}, /*stats*/
            std::nullopt /*distinctValueCount*/,
            0 /* nullCount */,
            false /* sorted */,
            sizeof(uint64_t) * sequenceCount /* represented_bytes */);

        return std::make_unique<SequenceVector<uint64_t> >(
            BaseVector::pool_,
            BaseVector::length_,
            std::move(hashValues),
            sequenceLengths_,
            SimpleVectorStats<uint64_t>{}, /*stats*/
            std::nullopt /*distinctCount*/,
            0 /* nullCount */,
            false /* sorted */,
            sizeof(uint64_t) * BaseVector::length_ /* represented_bytes */,
            0 /* nullSequenceCount */);
    }

    template<typename T>
    ksimd::batch<T> SequenceVector<T>::loadSIMDValueBufferAt(
        size_t byteOffset) const {
        if constexpr (std::is_same_v<T, bool>) {
            throw std::runtime_error(
                "Sequence encoding only supports SIMD operations on integers");
        } else {
            constexpr int kBatchSize = ksimd::batch<T>::size;
            auto startIndex = byteOffset / sizeof(T);
            if (checkLoadRange(startIndex, kBatchSize)) {
                return simd::setAll(valueAtFast(startIndex));
            }
            alignas(ksimd::default_arch::alignment()) T tmp[kBatchSize];
            for (int i = 0; i < kBatchSize; ++i) {
                tmp[i] = valueAtFast(startIndex + i);
            }
            return ksimd::load_aligned(tmp);
        }
    }

    template<typename T>
    bool SequenceVector<T>::checkLoadRange(size_t index, size_t count) const {
        // If the entire range is below the sequence threshold then we can load
        // everything at once
        offsetOfIndex(index); // set the internal index variables
        return (index + count) <= lastIndexRangeEnd_;
    }

    template<typename T>
    vector_size_t SequenceVector<T>::offsetOfIndex(vector_size_t index) const {
        POLLUX_DCHECK_LE(0, index);
        POLLUX_DCHECK_LT(index, BaseVector::length_);

        if (index < lastIndexRangeStart_) {
            // walk down the range indices until we are in the range requested
            do {
                --lastRangeIndex_;
                lastIndexRangeEnd_ = lastIndexRangeStart_;
                lastIndexRangeStart_ -= lengths_[lastRangeIndex_];
            } while (index < lastIndexRangeStart_);
        } else if (index >= lastIndexRangeEnd_) {
            // walk up the range indicies until we are in the range requested
            do {
                ++lastRangeIndex_;
                lastIndexRangeStart_ = lastIndexRangeEnd_;
                lastIndexRangeEnd_ += lengths_[lastRangeIndex_];
            } while (index >= lastIndexRangeEnd_);
        }

        return lastRangeIndex_;
    }

    static inline vector_size_t offsetOfIndex(
        const vector_size_t *lengths,
        vector_size_t index,
        vector_size_t *lastRangeBegin,
        vector_size_t *lastRangeEnd,
        vector_size_t *lastIndex) {
        if (index < *lastRangeBegin) {
            // walk down the range indicies until we are in the range requested
            do {
                --*lastIndex;
                *lastRangeEnd = *lastRangeBegin;
                *lastRangeBegin -= lengths[*lastIndex];
            } while (index < *lastRangeBegin);
        } else if (index >= *lastRangeEnd) {
            // walk up the range indicies until we are in the range requested
            do {
                ++*lastIndex;
                *lastRangeBegin = *lastRangeEnd;
                *lastRangeEnd += lengths[*lastIndex];
            } while (index >= *lastRangeEnd);
        }
        return *lastIndex;
    }
} // namespace kumo::pollux
